{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SageMaker LLM-as-Judge Evaluation - Basic Usage\n",
    "\n",
    "This notebook demonstrates the basic user-facing flow for creating and managing LLM-as-Judge evaluation jobs using the LLMAsJudgeEvaluator."
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Configure AWS credentials and region\n",
    "#! ada credentials update --provider=isengard --account=<> --role=Admin --profile=default --once\n",
    "#! aws configure set region us-west-2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Configuration\n",
    "REGION = 'us-west-2'\n",
    "S3_BUCKET = 's3://mufi-test-serverless-smtj/eval/'\n",
    "# DATASET = 'arn:aws:sagemaker:us-west-2:<>:hub-content/AIRegistry/DataSet/gen-qa-test-content/1.0.1'  # Dataset ARN or S3 URI\n",
    "DATASET = \"s3://my-sagemaker-sherpa-dataset/dataset/gen-qa-formatted-dataset/gen_qa.jsonl\"\n",
    "MLFLOW_ARN = 'arn:aws:sagemaker:us-west-2:<>:mlflow-tracking-server/mmlu-eval-experiment'"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Import Required Libraries\n",
    "\n",
    "Import the LLMAsJudgeEvaluator class."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import json\n",
    "from sagemaker.train.evaluate import LLMAsJudgeEvaluator\n",
    "from rich.pretty import pprint\n",
    "\n",
    "# Configure logging to show INFO messages\n",
    "import logging\n",
    "logging.basicConfig(\n",
    "    level=logging.INFO,\n",
    "    format='%(levelname)s - %(name)s - %(message)s'\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Create LLMAsJudgeEvaluator\n",
    "\n",
    "Create an LLMAsJudgeEvaluator instance with the desired evaluator model, dataset, and metrics.\n",
    "\n",
    "### Key Parameters:\n",
    "- `model`: Model Package (or Base Model) to be evaluated (required)\n",
    "- `evaluator_model`: Bedrock model ID to use as judge (required)\n",
    "- `dataset`: S3 URI or Dataset ARN (required)\n",
    "- `builtin_metrics`: List of built-in metrics (optional, no 'Builtin.' prefix needed)\n",
    "- `custom_metrics`: JSON string of custom metrics (optional)\n",
    "- `evaluate_base_model`: Whether to evaluate base model in addition to custom model (optional, default=True)\n",
    "- `mlflow_resource_arn`: MLflow tracking server ARN (optional)\n",
    "- `model_package_group`: Model package group ARN (optional)\n",
    "- `s3_output_path`: S3 output location (required)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Using custom metrics (as JSON string)\n",
    "\n",
    "Custom metrics must be provided as a properly escaped JSON string. You can either:\n",
    "1. Create a Python dict and use `json.dumps()` to convert it\n",
    "2. Provide a pre-escaped JSON string directly"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Method 1: Create dict and convert to JSON string\n",
    "custom_metric_dict = {\n",
    "    \"customMetricDefinition\": {\n",
    "        \"name\": \"PositiveSentiment\",\n",
    "        \"instructions\": (\n",
    "            \"You are an expert evaluator. Your task is to assess if the sentiment of the response is positive. \"\n",
    "            \"Rate the response based on whether it conveys positive sentiment, helpfulness, and constructive tone.\\n\\n\"\n",
    "            \"Consider the following:\\n\"\n",
    "            \"- Does the response have a positive, encouraging tone?\\n\"\n",
    "            \"- Is the response helpful and constructive?\\n\"\n",
    "            \"- Does it avoid negative language or criticism?\\n\\n\"\n",
    "            \"Rate on this scale:\\n\"\n",
    "            \"- Good: Response has positive sentiment\\n\"\n",
    "            \"- Poor: Response lacks positive sentiment\\n\\n\"\n",
    "            \"Here is the actual task:\\n\"\n",
    "            \"Prompt: {{prompt}}\\n\"\n",
    "            \"Response: {{prediction}}\"\n",
    "        ),\n",
    "        \"ratingScale\": [\n",
    "            {\"definition\": \"Good\", \"value\": {\"floatValue\": 1}},\n",
    "            {\"definition\": \"Poor\", \"value\": {\"floatValue\": 0}}\n",
    "        ]\n",
    "    }\n",
    "}\n",
    "\n",
    "# Convert to JSON string\n",
    "custom_metrics_json = json.dumps([custom_metric_dict])  # Note: wrap in list"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "\n",
    "# Create evaluator with custom metrics\n",
    "evaluator = LLMAsJudgeEvaluator(\n",
    "    # base_model='arn:aws:sagemaker:us-west-2:<>:model-package/Demo-test-deb-2/1',  # Required\n",
    "    model=\"arn:aws:sagemaker:us-west-2:<>:model-package/test-finetuned-models-gamma/28\",\n",
    "    evaluator_model=\"anthropic.claude-3-5-haiku-20241022-v1:0\",  # Required\n",
    "    dataset=DATASET,  # Required: S3 URI or Dataset ARN\n",
    "    builtin_metrics=[\"Completeness\", \"Faithfulness\"],  # Optional: Can combine with custom metrics\n",
    "    custom_metrics=custom_metrics_json,  # Optional: JSON string of custom metrics\n",
    "    mlflow_resource_arn=MLFLOW_ARN,  # Optional\n",
    "    # model_package_group=MODEL_PACKAGE_GROUP,  # Optional if BASE_MODEL is a Model Package ARN/Object\n",
    "    s3_output_path=S3_BUCKET,  # Required\n",
    "    evaluate_base_model=False\n",
    ")\n",
    "\n",
    "pprint(evaluator)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Optional] Example with multiple custom metrics"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# # Create multiple custom metrics\n",
    "# custom_metrics_list = [\n",
    "#     {\n",
    "#         \"customMetricDefinition\": {\n",
    "#             \"name\": \"GoodMetric\",\n",
    "#             \"instructions\": (\n",
    "#                 \"Assess if the response has positive sentiment. \"\n",
    "#                 \"Prompt: {{prompt}}\\nResponse: {{prediction}}\"\n",
    "#             ),\n",
    "#             \"ratingScale\": [\n",
    "#                 {\"definition\": \"Good\", \"value\": {\"floatValue\": 1}},\n",
    "#                 {\"definition\": \"Poor\", \"value\": {\"floatValue\": 0}}\n",
    "#             ]\n",
    "#         }\n",
    "#     },\n",
    "#     {\n",
    "#         \"customMetricDefinition\": {\n",
    "#             \"name\": \"BadMetric\",\n",
    "#             \"instructions\": (\n",
    "#                 \"Assess if the response has negative sentiment. \"\n",
    "#                 \"Prompt: {{prompt}}\\nResponse: {{prediction}}\"\n",
    "#             ),\n",
    "#             \"ratingScale\": [\n",
    "#                 {\"definition\": \"Bad\", \"value\": {\"floatValue\": 1}},\n",
    "#                 {\"definition\": \"Good\", \"value\": {\"floatValue\": 0}}\n",
    "#             ]\n",
    "#         }\n",
    "#     }\n",
    "# ]\n",
    "\n",
    "# # Convert list to JSON string\n",
    "# custom_metrics_json = json.dumps(custom_metrics_list)\n",
    "\n",
    "# # Create evaluator\n",
    "# evaluator = LLMAsJudgeEvaluator(\n",
    "#     base_model=BASE_MODEL,\n",
    "#     evaluator_model=\"anthropic.claude-3-5-haiku-20241022-v1:0\",\n",
    "#     dataset=DATASET,\n",
    "#     custom_metrics=custom_metrics_json,  # Multiple custom metrics\n",
    "#     s3_output_path=S3_BUCKET,\n",
    "# )\n",
    "\n",
    "# print(f\"✅ Created evaluator with {len(json.loads(custom_metrics_json))} custom metrics\")\n",
    "# pprint(evaluator)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Optional] Skipping base model evaluation (evaluate custom model only)\n",
    "\n",
    "By default, LLM-as-Judge evaluates both the base model and custom model. You can skip base model evaluation to save time and cost by setting `evaluate_base_model=False`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# # Define custom metrics (same as test script)\n",
    "# custom_metrics = \"[{\\\"customMetricDefinition\\\":{\\\"name\\\":\\\"GoodMetric\\\",\\\"instructions\\\":\\\"You are an expert evaluator. Your task is to assess if the sentiment of the response is positive. Rate the response based on whether it conveys positive sentiment, helpfulness, and constructive tone.\\\\n\\\\nConsider the following:\\\\n- Does the response have a positive, encouraging tone?\\\\n- Is the response helpful and constructive?\\\\n- Does it avoid negative language or criticism?\\\\n\\\\nRate on this scale:\\\\n- Good: Response has positive sentiment\\\\n- Poor: Response lacks positive sentiment\\\\n\\\\nHere is the actual task:\\\\nPrompt: {{prompt}}\\\\nResponse: {{prediction}}\\\",\\\"ratingScale\\\":[{\\\"definition\\\":\\\"Good\\\",\\\"value\\\":{\\\"floatValue\\\":1}},{\\\"definition\\\":\\\"Poor\\\",\\\"value\\\":{\\\"floatValue\\\":0}}]}},{\\\"customMetricDefinition\\\":{\\\"name\\\":\\\"BadMetric\\\",\\\"instructions\\\":\\\"You are an expert evaluator. Your task is to assess if the sentiment of the response is negative. Rate the response based on whether it conveys negative sentiment, unhelpfulness, or destructive tone.\\\\n\\\\nConsider the following:\\\\n- Does the response have a negative, discouraging tone?\\\\n- Is the response unhelpful or destructive?\\\\n- Does it use negative language or harsh criticism?\\\\n\\\\nRate on this scale:\\\\n- Bad: Response has negative sentiment\\\\n- Good: Response lacks negative sentiment\\\\n\\\\nHere is the actual task:\\\\nPrompt: {{prompt}}\\\\nResponse: {{prediction}}\\\",\\\"ratingScale\\\":[{\\\"definition\\\":\\\"Bad\\\",\\\"value\\\":{\\\"floatValue\\\":1}},{\\\"definition\\\":\\\"Good\\\",\\\"value\\\":{\\\"floatValue\\\":0}}]}}]\"\n",
    "\n",
    "# # Create evaluator that only evaluates the custom model (matching test script exactly)\n",
    "# evaluator = LLMAsJudgeEvaluator(\n",
    "#     base_model=BASE_MODEL,\n",
    "#     evaluator_model=\"anthropic.claude-3-5-haiku-20241022-v1:0\",\n",
    "#     dataset=DATASET,\n",
    "#     builtin_metrics=[\"Completeness\", \"Faithfulness\", \"Helpfulness\"],\n",
    "#     custom_metrics=custom_metrics,\n",
    "#     mlflow_resource_arn=MLFLOW_ARN,\n",
    "#     model_package_group=MODEL_PACKAGE_GROUP,\n",
    "#     model_artifact=MODEL_ARTIFACT,\n",
    "#     s3_output_path=S3_BUCKET,\n",
    "#     evaluate_base_model=False,  # KEY: Skip base model evaluation\n",
    "# )\n",
    "\n",
    "# print(\"✅ Created evaluator (custom model only)\")\n",
    "# pprint(evaluator)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Run LLM-as-Judge Evaluation\n",
    "\n",
    "Start the evaluation job. The evaluator will:\n",
    "1. Generate inference responses from the base model (if evaluate_base_model=True)\n",
    "2. Generate inference responses from the custom model\n",
    "3. Use the judge model to evaluate responses with built-in and custom metrics"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Run evaluation\n",
    "execution = evaluator.evaluate()\n",
    "\n",
    "print(f\"✅ Evaluation job started!\")\n",
    "print(f\"Job ARN: {execution.arn}\")\n",
    "print(f\"Job Name: {execution.name}\")\n",
    "print(f\"Status: {execution.status.overall_status}\")\n",
    "\n",
    "pprint(execution)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Check Job Status\n",
    "\n",
    "Refresh and display the current job status with step details."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Refresh status\n",
    "execution.refresh()\n",
    "\n",
    "# Display job status using rich pprint\n",
    "pprint(execution.status)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Monitor Pipeline Execution\n",
    "\n",
    "Poll the pipeline status until it reaches a terminal state (Succeeded, Failed, or Stopped)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Wait for job completion (optional)\n",
    "# This will poll every 5 seconds for up to 1 hour\n",
    "execution.wait(poll=5, timeout=3600)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Display results\n",
    "execution.show_results(limit=10, offset=0, show_explanations=False)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieve an Existing Job\n",
    "\n",
    "You can retrieve and inspect any existing evaluation job using its ARN."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Get an existing job by ARN\n",
    "# Replace with your actual pipeline execution ARN\n",
    "existing_arn = 'arn:aws:sagemaker:us-west-2:<>:pipeline/SagemakerEvaluation-llmasjudge/execution/4hr7446yft1d'  # or use a specific ARN\n",
    "\n",
    "from sagemaker.train.evaluate import EvaluationPipelineExecution\n",
    "from rich.pretty import pprint\n",
    "\n",
    "existing_execution = EvaluationPipelineExecution.get(\n",
    "    arn=existing_arn,\n",
    "    region=\"us-west-2\"\n",
    ")\n",
    "pprint(existing_execution.status)\n",
    "\n",
    "existing_execution.show_results(limit=5, offset=0, show_explanations=False)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get All LLM-as-Judge Evaluations\n",
    "\n",
    "Retrieve all LLM-as-Judge evaluation jobs."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from sagemaker.train.evaluate import LLMAsJudgeEvaluator\n",
    "\n",
    "# Get all LLM-as-Judge evaluations as an iterator\n",
    "all_executions = list(LLMAsJudgeEvaluator.get_all(region=\"us-west-2\"))\n",
    "\n",
    "print(f\"Found {len(all_executions)} LLM-as-Judge evaluation jobs\")\n",
    "for execution in all_executions:\n",
    "    print(f\"  - {execution.name}: {execution.status.overall_status}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stop a Running Job (Optional)\n",
    "\n",
    "If needed, you can stop a running evaluation job."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Uncomment to stop the job\n",
    "# execution.stop()\n",
    "# print(f\"Execution stopped. Status: {execution.status.overall_status}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset Support\n",
    "\n",
    "The `dataset` parameter supports two formats:\n",
    "\n",
    "### 1. S3 URI\n",
    "```python\n",
    "dataset=\"s3://my-bucket/path/to/dataset.jsonl\"\n",
    "```\n",
    "\n",
    "### 2. Dataset ARN (AI Registry)\n",
    "```python\n",
    "dataset=\"arn:aws:sagemaker:us-west-2:123456789012:hub-content/AIRegistry/DataSet/my-dataset/1.0.0\"\n",
    "```\n",
    "\n",
    "The evaluator automatically detects which format is provided and uses the appropriate data source configuration."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
